package org.nerve.webmagic.downloader;

import org.nerve.webmagic.PhantomJsDriverPool;
import org.nerve.webmagic.utils.ScreenShotTool;
import org.openqa.selenium.By;
import org.openqa.selenium.WebDriver;
import org.openqa.selenium.WebElement;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import us.codecraft.webmagic.Page;
import us.codecraft.webmagic.Request;
import us.codecraft.webmagic.Site;
import us.codecraft.webmagic.Task;
import us.codecraft.webmagic.downloader.Downloader;
import us.codecraft.webmagic.selector.PlainText;

import java.io.Closeable;
import java.io.IOException;

import static org.nerve.webmagic.bean.WebPage.SCREENSHOT;

/**
 * org.nerve.webmagic.spider
 * Created by zengxm on 2016/4/20 0020.
 */
public class PhantomJSDownloader2 implements Downloader, Closeable{
	private static Logger logger= LoggerFactory.getLogger(PhantomJSDownloader2.class);
	private PhantomJsDriverPool driverPool;

	private ThreadLocal<ScreenShotTool> screenThreadLocal=new ThreadLocal<>();

	private int sleepTime = 0;      //等待phantomjs加载网页的时间
	private int poolSize = 3;
	private boolean takeScreenShot=false;   //是否截图

	public PhantomJSDownloader2(){
		this(1,500);
	}
	public PhantomJSDownloader2(int poolSize){
		this(poolSize, 500);
	}
	public PhantomJSDownloader2(int poolSize, int sleepTime){
		this.poolSize=poolSize;
		this.sleepTime=sleepTime;

		init();
	}

	private void init(){
		if(driverPool==null){
			synchronized (this){
				driverPool=new PhantomJsDriverPool(poolSize);
			}
		}
	}

	public PhantomJSDownloader2 setSleepTime(int sleepTime) {
		this.sleepTime = sleepTime;
		return this;
	}

	/**
	 * 设置是否截图
	 * 如果为true，则将截图赋值到 page 对象的 SCREENSHOT 上。
	 * @param takeScreenShot
	 * @return
	 */
	public PhantomJSDownloader2 setTakeScreenShot(boolean takeScreenShot) {
		this.takeScreenShot = takeScreenShot;
		return this;
	}

	@Override
	public Page download(Request request, Task task) {
		final Site site=task.getSite();
		final WebDriver webDriver;
		try {
			webDriver = driverPool.get(site);
		} catch (InterruptedException e) {
			logger.warn("interrupted", e);
			return null;
		}
		logger.info("[WebDriver={}]downloading page {}, thread={}",
				webDriver.toString(),
				request.getUrl(),
				Thread.currentThread().getName());

		//加载网页
		webDriver.get(request.getUrl());
		try {
			Thread.sleep(sleepTime);
		} catch (InterruptedException e) {
			logger.error("try to sleep {} ms, but error on:{}", sleepTime, e.getMessage());
		}

		WebElement webElement = webDriver.findElement(By.xpath("/html"));
		String content = webElement.getAttribute("outerHTML");
		Page page = new Page();
		page.setRawText(content);
		page.setUrl(new PlainText(request.getUrl()));
		page.setRequest(request);

		if(takeScreenShot){
			page.putField(SCREENSHOT, taskScreenshot(webDriver));
		}

		//用完后，返还webDriver
		driverPool.returnToPool(webDriver);
		return page;
	}

	private String taskScreenshot(final WebDriver driver){
		ScreenShotTool tool=screenThreadLocal.get();
		if(tool==null){
			tool=new ScreenShotTool();
			screenThreadLocal.set(tool);
		}

		return tool.taskScreenshot(driver);
	}

	@Override
	public void setThread(int i) {
		this.poolSize=i;
	}

	@Override
	public void close() throws IOException {
		if(driverPool!=null)
			driverPool.closeAll();
	}
}
